import torch
from torch.cuda.amp import GradScaler, autocast
from torchattacks.attack import Attack
from attack.libfastattack.lossfunc import *


class FastFGSMTrain(Attack):

    def __init__(self, model, eps, cls_num_list=None,mode="Normal", weights=False):
        super().__init__("FastFGSMTrain", model)
        self.eps = eps
        self._supported_mode = ['default', 'targeted']
        self.scaler = GradScaler()
        self.alpha = 1.25
        self.cls_num_list = cls_num_list
        self.mode = mode
        self.weights = weights

    def forward(self, images, labels):

        mode = self.mode
        images = images.clone().detach().to(self.device)
        labels = labels.clone().detach().to(self.device)

        adv_images = images.clone().detach()

        # Starting at a uniformly random point
        adv_images = adv_images + torch.empty_like(adv_images).uniform_(-self.eps, self.eps)

        if self._targeted:
            target_labels = self._get_target_label(images, labels)
        if mode == "Normal":
            loss = torch.nn.CrossEntropyLoss()
        elif mode == "CalFAT":
            loss = nn.KLDivLoss(size_average=False)
            marloss = maxMarginLoss_kl(cls_num_list=self.cls_num_list, s=10,weight=None).cuda()
        else:
            assert 1 == 0

        adv_images.requires_grad = True

        # Accelarating forward propagation
        with autocast():
            outputs = self.model(adv_images)
            origin_outputs = self.model(images)
            
            if mode == "Normal":
                # Calculate loss
                if self._targeted:
                    cost = -loss(outputs, target_labels)
                else:
                    cost = loss(outputs, labels)
            elif mode == "CalFAT":
                cost = loss(
                        F.log_softmax(marloss(outputs), dim=1),
                        F.softmax(marloss(origin_outputs), dim =1)
                    )
            else:
                assert 1 == 2

        # Accelerating Gradient
        scaled_loss = self.scaler.scale(cost)
        # Update adversarial images
        grad = torch.autograd.grad(scaled_loss, adv_images,
                                   retain_graph=False, create_graph=False)[0]

        adv_images_ = adv_images.detach() + self.alpha * self.eps*grad.sign()
        delta = torch.clamp(adv_images_ - images, min=-self.eps, max=self.eps)
        different_rounds = torch.ones(images.shape[0], dtype=torch.int).cuda()
        weights = different_rounds.float() / different_rounds.float().sum().float()
        if self.weights:
            return torch.clamp(images + delta, min=0, max=1).detach(), weights,different_rounds
        else:
            return torch.clamp(images + delta, min=0, max=1).detach()